/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 性能评估器相关头文件
 * Author: qingjicheng
 * Create: 2022-8-1
 */

#ifndef CONNECTED_GRAPH_ANALYZER_H
#define CONNECTED_GRAPH_ANALYZER_H
#include <map>
#include <functional>
#include "graph_common_def.h"
#include "connected_graph.h"

namespace HcclEvaluator {
struct AlgSplitParams {
    const ConnectedGraph &graph;
    HcclOpInfo opInfo;
};

class ConnectedGraphAnalyzer {
public:
    explicit ConnectedGraphAnalyzer();
    virtual ~ConnectedGraphAnalyzer();
    HcclResult AnalyzeGraphTimeCost(const ConnectedGraph &graph, HcclOpInfo &op, float &cost);
private:
    std::vector<std::vector<SubgraphCalcInfo>> AlgSplitAllToAllV(AlgSplitParams para);
    std::vector<std::vector<SubgraphCalcInfo>> AlgSplitAllGather(AlgSplitParams para);
    std::vector<std::vector<SubgraphCalcInfo>> AlgSplitAllReduce(AlgSplitParams para);
    std::vector<std::vector<SubgraphCalcInfo>> AlgSplitBroadcast(AlgSplitParams para);
    std::vector<std::vector<SubgraphCalcInfo>> AlgSplitReduceScatter(AlgSplitParams para);
    std::vector<std::vector<SubgraphCalcInfo>> AlgSplitReduce(AlgSplitParams para);
private:
    HcclResult CheckOpInfo(HcclOpInfo &op) const;
    float CalcLevelCost(std::vector<SubgraphCalcInfo> &calcInfos,
        const std::vector<std::unique_ptr<ConnectedSubgraph>> &subgraphs);
    float CalcSubgraphCost(std::vector<SubgraphCalcInfo> &calcInfos,
        const std::unique_ptr<ConnectedSubgraph>& subgraph);
    template<typename ArrType> ArrType GetArrMaxValue(std::vector<ArrType> arr)
    {
        ArrType maxValue = {0};
        for (ArrType &val : arr) {
            if (val > maxValue) {
                maxValue = val;
            }
        }
        return maxValue;
    }
private:
    std::map<HcclOpType, std::function<std::vector<std::vector<SubgraphCalcInfo>>(AlgSplitParams)>> opFuncTable_;
};
}
#endif